## 用于生成jsMind的json数据
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
import re
from typing import Dict
import json
from PyPDF2 import PdfReader
import uuid  # 在文件顶部添加此导入

class MindModel:
    def __init__(self,api_key,api_base,model_name="glm-4"):
        self.llm = ChatOpenAI(
            temperature=0.95,
            model=model_name,
            openai_api_key=api_key,
            openai_api_base=api_base
        )
        self.set_prompt()

    def set_prompt(self):
        self.system_template = """
        你是一个专业文档转换助手，请将PDF内容转换为层级清晰的思维导图结构，要求如下：

        1. **输出格式要求**：
        - 只输出标题结构，不输出详细段落内容
        - 使用Markdown标题格式（# → ###），最多三级标题
        - 层级结构清晰，逻辑合理，不超过三级嵌套

        2. **内容处理规则**：
        - 将没有标题的段落内容总结为对应层级的标题信息
        - 相似/重复的内容合并为一个标题
        - 将详细描述和例子转化为简洁的标题点
        - 确保标题间有明确的层级关系和逻辑连接

        3. **标题处理规则**：
        - 标题应简洁明了，通常5-10个字为宜
        - 同级标题应表达同等重要性的内容
        - 下级标题应是上级标题的细化或补充
        - 保持原文的知识结构和关键术语

        4. **其他要求**：
        - 不输出数学公式、表格、图片和代码块
        - 不添加注释或解释性文字
        - 保持中英文原文语种
        """
        self.prompt = ChatPromptTemplate(
            messages=[
                SystemMessagePromptTemplate.from_template(
                    self.system_template,
                ),
                HumanMessagePromptTemplate.from_template("文档内容：{content}")
            ]
        )
    
    def generate_md(self,content):
        response = self.llm.invoke(self.prompt.format(content=content))
        return response.content
    
class MarkdownToJsMind:
    def __init__(self):
        pass  # 移除 self.node_id = 0，因为我们将使用 uuid

    def _gen_id(self) -> str:
        """生成唯一节点ID"""
        # 生成 UUID 并移除横线，截取前 16 位
        return str(uuid.uuid4()).replace('-', '')[:16]

    def _parse_header(self, line: str) -> Dict:
        """解析Markdown标题"""
        match = re.match(r"^(#+)\s+(.+)$", line)
        if match:
            return {
                "level": len(match.group(1)),
                "text": match.group(2).strip()
            }
        return None

    def build_tree(self, md_str: str) -> Dict:
        """
        将Markdown字符串转换为jsMind格式的JSON，并实现左右平衡布局
        参数：
            md_str: 完整的Markdown内容字符串
        返回：  
            符合jsMind规范的树形JSON
        """
        # 初始化数据结构
        mind_map = {
            "meta": {"name": "Auto-generated MindMap"},
            "format": "node_tree",
            "data": {"id": "root", "topic": "Root Node", "children": []}
        }

        # 按行处理Markdown内容  
        lines = md_str.splitlines()
        stack = [mind_map["data"]]  # 用栈来维护节点层级
        
        # 收集所有一级节点，以便后续分配方向
        first_level_nodes = []

        for line in lines:
            line = line.strip()
            if not line:
                continue

            # 处理标题
            header = self._parse_header(line)
            if header:
                level = header["level"]
                topic = header["text"]

                # 调整栈，使当前节点成为正确层级的子节点
                while len(stack) > level:
                    stack.pop()
                
                # 创建新节点并添加到父节点的children
                new_node = {"id": self._gen_id(), "topic": topic, "children": []}
                stack[-1]["children"].append(new_node)
                
                # 记录一级节点（level为1的节点）
                if level == 1:
                    first_level_nodes.append(new_node)
                
                stack.append(new_node)  # 将新节点压入栈
                continue

            # 处理列表项（支持无序列表）  
            if line.startswith("- "):
                list_item = line[2:].strip()
                stack[-1]["children"].append({
                    "id": self._gen_id(), 
                    "topic": list_item
                })
                continue

            # 处理普通文本段落
            stack[-1]["children"].append({
                "id": self._gen_id(),
                "topic": line
            })

        # 为一级节点分配左右方向，实现平衡布局
        for i, node in enumerate(first_level_nodes):
            # 偶数索引节点放右侧，奇数索引节点放左侧
            node["direction"] = "right" if i % 2 == 0 else "left"

        return mind_map
    
# ========================
# 3. 合并逻辑实现：递归合并两个 jsMind 树
# ========================
def merge_node_children(old_children, new_children):
    """
    合并两个节点列表。对每个 new_node，若在 old_children 中根据特定标识（page_number 或 topic）能找到匹配节点，
    则递归合并其 children；否则直接追加。
    """
    for new_node in new_children:
        merged = False
        # 如果 new_node 带有 page_number，则以 page_number 为匹配依据
        if "page_number" in new_node:
            for oc in old_children:
                if oc.get("page_number") == new_node.get("page_number"):
                    # 递归合并容器节点下的 children
                    oc["children"] = merge_node_children(oc.get("children", []), new_node.get("children", []))
                    merged = True
                    break
        else:
            # 否则以 topic 作为匹配依据（也可扩展其他标识，如节点 id 或哈希值）
            for oc in old_children:
                if oc.get("topic") == new_node.get("topic"):
                    # 如果双方都有 children，则递归合并
                    if "children" in new_node and new_node["children"]:
                        oc["children"] = merge_node_children(oc.get("children", []), new_node["children"])
                    merged = True
                    break
        if not merged:
            old_children.append(new_node)
    return old_children

def merge_trees(old_tree: dict, new_tree: dict) -> dict:
    """
    将新生成的增量树 new_tree 合并到已有树 old_tree 中。
    思路：直接将 new_tree["data"]["children"] 与 old_tree["data"]["children"] 递归合并。
    合并后重新分配一级节点的左右方向，确保平衡布局。
    """
    old_children = old_tree["data"].get("children", [])
    new_children = new_tree["data"].get("children", [])
    merged_children = merge_node_children(old_children, new_children)
    old_tree["data"]["children"] = merged_children
    
    # 合并后重新分配一级节点的左右方向
    first_level_nodes = old_tree["data"]["children"]
    for i, node in enumerate(first_level_nodes):
        # 偶数索引节点放右侧，奇数索引节点放左侧
        node["direction"] = "right" if i % 2 == 0 else "left"
    
    return old_tree

    
if __name__ == "__main__":
    api_key = "sk-51dd4f350cdd49f58083396702d29b25"
    api_base = "https://api.deepseek.com"
    model_name = "deepseek-chat"
    mind_model = MindModel(api_key,api_base,model_name)
    mind_model.set_prompt()
    ##使用pdfread读取server\static\docs\CISSP官方学习手册（第9版） ([美]迈克·查普尔(Mike Chapple)) (Z-Library).pdf文件的  10-15页和20-25页内容
    pdf_reader = PdfReader("E:\\APP\\QuickLearn\\server\\static\\docs\\CISSP官方学习手册（第9版） ([美]迈克·查普尔(Mike Chapple)) (Z-Library).pdf")
    content_1 = ""
    content_2 = ""
    for page in pdf_reader.pages[52:57]:
        content_1 += page.extract_text()
    for page in pdf_reader.pages[84:89]:
        content_2 += page.extract_text()

   


    md_1 = mind_model.generate_md(content_1)

    markdown_to_js_mind = MarkdownToJsMind()
    js_mind_1 = markdown_to_js_mind.build_tree(md_1)
    print("第一部分的：",json.dumps(js_mind_1, ensure_ascii=False))

    print("\n============================================================\n")
    md_2 = mind_model.generate_md(content_2)
    js_mind_2 = markdown_to_js_mind.build_tree(md_2)
    print("第二部分的：",json.dumps(js_mind_2, ensure_ascii=False))

    print("\n============================================================\n")

    # 合并两部分
    merged_tree = merge_trees(js_mind_1, js_mind_2)
    print("合并后的树：", json.dumps(merged_tree, ensure_ascii=False))
